#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
攻击数据验证系统
验证攻击链测试数据的真实性和完整性
"""

import os
import sys
import json
import hashlib
import time
from datetime import datetime, timedelta
from pathlib import Path
import logging

class AttackDataValidator:
    def __init__(self, evidence_file):
        self.evidence_file = Path(evidence_file)
        self.validation_results = {
            "validation_id": f"VALIDATION_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            "evidence_file": str(evidence_file),
            "validation_time": datetime.now().isoformat(),
            "checks_performed": [],
            "validation_score": 0,
            "authenticity_indicators": [],
            "anomalies_detected": [],
            "overall_status": "PENDING"
        }
        
        self.setup_logging()
        self.load_evidence_data()
    
    def setup_logging(self):
        """设置日志记录"""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s'
        )
        self.logger = logging.getLogger(__name__)
    
    def load_evidence_data(self):
        """加载证据数据"""
        try:
            with open(self.evidence_file, 'r', encoding='utf-8') as f:
                self.evidence_data = json.load(f)
            self.logger.info(f"成功加载证据文件: {self.evidence_file}")
        except Exception as e:
            self.logger.error(f"加载证据文件失败: {e}")
            raise
    
    def validate_metadata_consistency(self):
        """验证元数据一致性"""
        check_name = "metadata_consistency"
        self.logger.info("开始验证元数据一致性...")
        
        issues = []
        score = 0
        
        # 检查必要字段
        required_fields = ["test_id", "start_time", "end_time", "total_duration"]
        for field in required_fields:
            if field in self.evidence_data["test_metadata"]:
                score += 10
            else:
                issues.append(f"缺少必要字段: {field}")
        
        # 检查时间一致性
        try:
            start_time = datetime.fromisoformat(self.evidence_data["test_metadata"]["start_time"])
            end_time = datetime.fromisoformat(self.evidence_data["test_metadata"]["end_time"])
            
            if end_time > start_time:
                score += 20
                self.validation_results["authenticity_indicators"].append(
                    "时间戳逻辑正确: 结束时间晚于开始时间"
                )
            else:
                issues.append("时间戳逻辑错误: 结束时间早于开始时间")
        except Exception as e:
            issues.append(f"时间戳格式错误: {e}")
        
        # 检查测试ID格式
        test_id = self.evidence_data["test_metadata"]["test_id"]
        if test_id.startswith("ATTACK_TEST_") and len(test_id) == 28:
            score += 15
            self.validation_results["authenticity_indicators"].append(
                "测试ID格式正确"
            )
        else:
            issues.append("测试ID格式不正确")
        
        self.validation_results["checks_performed"].append({
            "check": check_name,
            "score": score,
            "max_score": 45,
            "issues": issues,
            "status": "PASS" if score >= 35 else "FAIL"
        })
        
        return score >= 35
    
    def validate_attack_timeline(self):
        """验证攻击时间线"""
        check_name = "attack_timeline"
        self.logger.info("开始验证攻击时间线...")
        
        issues = []
        score = 0
        
        timeline = self.evidence_data.get("attack_timeline", [])
        
        # 检查阶段完整性
        expected_phases = [
            "reconnaissance", "initial_access", "privilege_escalation",
            "persistence", "lateral_movement", "data_collection",
            "exfiltration", "emergency_response"
        ]
        
        found_phases = [phase["phase"] for phase in timeline]
        
        for phase in expected_phases:
            if phase in found_phases:
                score += 5
            else:
                issues.append(f"缺少攻击阶段: {phase}")
        
        # 检查时间顺序
        if len(timeline) > 1:
            time_ordered = True
            for i in range(1, len(timeline)):
                prev_time = datetime.fromisoformat(timeline[i-1]["start_time"])
                curr_time = datetime.fromisoformat(timeline[i]["start_time"])
                if curr_time < prev_time:
                    time_ordered = False
                    break
            
            if time_ordered:
                score += 20
                self.validation_results["authenticity_indicators"].append(
                    "攻击阶段时间顺序正确"
                )
            else:
                issues.append("攻击阶段时间顺序错误")
        
        # 检查每个阶段的数据完整性
        for phase in timeline:
            if all(key in phase for key in ["phase", "start_time", "end_time", "status", "results"]):
                score += 2
            else:
                issues.append(f"阶段 {phase.get('phase', 'unknown')} 数据不完整")
        
        self.validation_results["checks_performed"].append({
            "check": check_name,
            "score": score,
            "max_score": 56,
            "issues": issues,
            "status": "PASS" if score >= 45 else "FAIL"
        })
        
        return score >= 45
    
    def validate_network_traffic(self):
        """验证网络流量数据"""
        check_name = "network_traffic"
        self.logger.info("开始验证网络流量数据...")
        
        issues = []
        score = 0
        
        traffic = self.evidence_data.get("network_traffic", [])
        
        if len(traffic) > 0:
            score += 20
            self.validation_results["authenticity_indicators"].append(
                f"包含 {len(traffic)} 个网络数据包"
            )
        else:
            issues.append("缺少网络流量数据")
        
        # 检查数据包字段完整性
        required_fields = ["timestamp", "src_ip", "dst_ip", "protocol", "dst_port"]
        valid_packets = 0
        
        for packet in traffic[:10]:  # 检查前10个数据包
            if all(field in packet for field in required_fields):
                valid_packets += 1
        
        if valid_packets >= 8:
            score += 15
        elif valid_packets >= 5:
            score += 10
        else:
            issues.append("网络数据包字段不完整")
        
        # 检查IP地址格式
        valid_ips = 0
        for packet in traffic[:5]:
            src_ip = packet.get("src_ip", "")
            dst_ip = packet.get("dst_ip", "")
            if self.is_valid_ip(src_ip) and self.is_valid_ip(dst_ip):
                valid_ips += 1
        
        if valid_ips >= 4:
            score += 15
            self.validation_results["authenticity_indicators"].append(
                "IP地址格式正确"
            )
        else:
            issues.append("IP地址格式错误")
        
        # 检查攻击阶段分布
        phases = set()
        for packet in traffic:
            if "phase" in packet:
                phases.add(packet["phase"])
        
        if len(phases) >= 3:
            score += 10
            self.validation_results["authenticity_indicators"].append(
                f"网络流量覆盖 {len(phases)} 个攻击阶段"
            )
        
        self.validation_results["checks_performed"].append({
            "check": check_name,
            "score": score,
            "max_score": 60,
            "issues": issues,
            "status": "PASS" if score >= 45 else "FAIL"
        })
        
        return score >= 45
    
    def validate_system_changes(self):
        """验证系统变更数据"""
        check_name = "system_changes"
        self.logger.info("开始验证系统变更数据...")
        
        issues = []
        score = 0
        
        changes = self.evidence_data.get("system_changes", [])
        
        if len(changes) > 0:
            score += 15
        else:
            issues.append("缺少系统变更数据")
        
        # 检查变更类型多样性
        change_types = set()
        for change in changes:
            if "type" in change:
                change_types.add(change["type"])
        
        expected_types = ["file_creation", "process_creation", "network_connection"]
        found_types = len(change_types.intersection(expected_types))
        
        score += found_types * 10
        
        if found_types >= 2:
            self.validation_results["authenticity_indicators"].append(
                f"包含 {found_types} 种系统变更类型"
            )
        
        # 检查文件路径真实性
        file_changes = [c for c in changes if c.get("type") == "file_creation"]
        realistic_paths = 0
        
        for change in file_changes:
            path = change.get("path", "")
            if any(prefix in path for prefix in ["/tmp/", "/var/", "/etc/", "/usr/", "/home/"]):
                realistic_paths += 1
        
        if realistic_paths >= len(file_changes) * 0.8:
            score += 15
            self.validation_results["authenticity_indicators"].append(
                "文件路径符合Linux系统规范"
            )
        
        # 检查哈希值格式
        hash_valid = 0
        for change in changes:
            if "hash" in change:
                hash_val = change["hash"]
                if len(hash_val) == 64:  # SHA256
                    hash_valid += 1
        
        if hash_valid > 0:
            score += 10
            self.validation_results["authenticity_indicators"].append(
                f"{hash_valid} 个文件包含有效哈希值"
            )
        
        self.validation_results["checks_performed"].append({
            "check": check_name,
            "score": score,
            "max_score": 70,
            "issues": issues,
            "status": "PASS" if score >= 50 else "FAIL"
        })
        
        return score >= 50
    
    def validate_persistence_mechanisms(self):
        """验证持久化机制"""
        check_name = "persistence_mechanisms"
        self.logger.info("开始验证持久化机制...")
        
        issues = []
        score = 0
        
        mechanisms = self.evidence_data.get("persistence_mechanisms", [])
        
        if len(mechanisms) > 0:
            score += 20
        else:
            issues.append("缺少持久化机制数据")
        
        # 检查持久化类型多样性
        mech_types = set()
        for mech in mechanisms:
            if "type" in mech:
                mech_types.add(mech["type"])
        
        expected_types = ["cron_job", "systemd_service", "ssh_key", "ld_preload"]
        found_types = len(mech_types.intersection(expected_types))
        
        score += found_types * 15
        
        if found_types >= 3:
            self.validation_results["authenticity_indicators"].append(
                f"使用了 {found_types} 种持久化技术"
            )
        
        # 检查隐蔽级别
        stealth_levels = [mech.get("stealth_level") for mech in mechanisms if "stealth_level" in mech]
        if stealth_levels:
            score += 10
            self.validation_results["authenticity_indicators"].append(
                "包含隐蔽级别评估"
            )
        
        self.validation_results["checks_performed"].append({
            "check": check_name,
            "score": score,
            "max_score": 90,
            "issues": issues,
            "status": "PASS" if score >= 60 else "FAIL"
        })
        
        return score >= 60
    
    def validate_data_exfiltration(self):
        """验证数据渗透"""
        check_name = "data_exfiltration"
        self.logger.info("开始验证数据渗透...")
        
        issues = []
        score = 0
        
        exfiltration = self.evidence_data.get("data_exfiltration", [])
        
        if len(exfiltration) > 0:
            score += 20
        else:
            issues.append("缺少数据渗透数据")
        
        # 检查渗透动作类型
        action_types = set()
        for action in exfiltration:
            if "action" in action:
                action_types.add(action["action"])
        
        expected_actions = ["file_collection", "database_dump", "network_transfer"]
        found_actions = len(action_types.intersection(expected_actions))
        
        score += found_actions * 15
        
        if found_actions >= 2:
            self.validation_results["authenticity_indicators"].append(
                f"执行了 {found_actions} 种数据渗透动作"
            )
        
        # 检查加密和压缩
        encrypted_actions = [a for a in exfiltration if "encryption" in a]
        if encrypted_actions:
            score += 15
            self.validation_results["authenticity_indicators"].append(
                "数据传输使用了加密"
            )
        
        self.validation_results["checks_performed"].append({
            "check": check_name,
            "score": score,
            "max_score": 65,
            "issues": issues,
            "status": "PASS" if score >= 45 else "FAIL"
        })
        
        return score >= 45
    
    def is_valid_ip(self, ip):
        """验证IP地址格式"""
        try:
            parts = ip.split('.')
            return len(parts) == 4 and all(0 <= int(part) <= 255 for part in parts)
        except:
            return False
    
    def calculate_overall_score(self):
        """计算总体验证分数"""
        total_score = 0
        max_total_score = 0
        
        for check in self.validation_results["checks_performed"]:
            total_score += check["score"]
            max_total_score += check["max_score"]
        
        if max_total_score > 0:
            percentage = (total_score / max_total_score) * 100
            self.validation_results["validation_score"] = round(percentage, 2)
        
        # 确定总体状态
        if self.validation_results["validation_score"] >= 85:
            self.validation_results["overall_status"] = "HIGHLY_AUTHENTIC"
        elif self.validation_results["validation_score"] >= 70:
            self.validation_results["overall_status"] = "AUTHENTIC"
        elif self.validation_results["validation_score"] >= 50:
            self.validation_results["overall_status"] = "PARTIALLY_AUTHENTIC"
        else:
            self.validation_results["overall_status"] = "QUESTIONABLE"
    
    def run_validation(self):
        """运行完整验证"""
        self.logger.info("开始运行攻击数据验证...")
        
        # 执行各项验证
        validations = [
            self.validate_metadata_consistency,
            self.validate_attack_timeline,
            self.validate_network_traffic,
            self.validate_system_changes,
            self.validate_persistence_mechanisms,
            self.validate_data_exfiltration
        ]
        
        for validation in validations:
            try:
                validation()
            except Exception as e:
                self.logger.error(f"验证过程中发生错误: {e}")
                self.validation_results["anomalies_detected"].append(f"验证错误: {e}")
        
        # 计算总体分数
        self.calculate_overall_score()
        
        self.logger.info(f"验证完成，总体分数: {self.validation_results['validation_score']}%")
        
        return self.validation_results
    
    def generate_validation_report(self):
        """生成验证报告"""
        report_dir = Path("reports/validation_reports")
        report_dir.mkdir(parents=True, exist_ok=True)
        
        # JSON报告
        json_file = report_dir / f"validation_{self.validation_results['validation_id']}.json"
        with open(json_file, 'w', encoding='utf-8') as f:
            json.dump(self.validation_results, f, indent=2, ensure_ascii=False)
        
        # Markdown报告
        md_file = report_dir / f"validation_{self.validation_results['validation_id']}.md"
        with open(md_file, 'w', encoding='utf-8') as f:
            f.write(f"# 攻击数据验证报告\n\n")
            f.write(f"**验证ID**: {self.validation_results['validation_id']}\n")
            f.write(f"**证据文件**: {self.validation_results['evidence_file']}\n")
            f.write(f"**验证时间**: {self.validation_results['validation_time']}\n")
            f.write(f"**验证分数**: {self.validation_results['validation_score']}%\n")
            f.write(f"**总体状态**: {self.validation_results['overall_status']}\n\n")
            
            f.write("## 验证检查结果\n\n")
            for check in self.validation_results["checks_performed"]:
                f.write(f"### {check['check'].title()}\n")
                f.write(f"- **分数**: {check['score']}/{check['max_score']}\n")
                f.write(f"- **状态**: {check['status']}\n")
                if check['issues']:
                    f.write(f"- **问题**: {', '.join(check['issues'])}\n")
                f.write("\n")
            
            f.write("## 真实性指标\n\n")
            for indicator in self.validation_results["authenticity_indicators"]:
                f.write(f"- {indicator}\n")
            
            if self.validation_results["anomalies_detected"]:
                f.write("\n## 检测到的异常\n\n")
                for anomaly in self.validation_results["anomalies_detected"]:
                    f.write(f"- {anomaly}\n")
        
        return {
            "json_report": str(json_file),
            "markdown_report": str(md_file)
        }

def main():
    """主函数"""
    
    if len(sys.argv) != 2:
        print("用法: python attack_validation_system.py <evidence_file>")
        sys.exit(1)
    
    evidence_file = sys.argv[1]
    
    if not Path(evidence_file).exists():
        print(f"错误: 证据文件不存在: {evidence_file}")
        sys.exit(1)
    
    print("=" * 60)
    print("攻击数据验证系统")
    print("=" * 60)
    
    try:
        # 创建验证器
        validator = AttackDataValidator(evidence_file)
        
        # 运行验证
        print("\n[+] 开始验证攻击数据...")
        results = validator.run_validation()
        
        # 生成报告
        print("\n[+] 生成验证报告...")
        reports = validator.generate_validation_report()
        
        # 输出结果
        print("\n" + "=" * 60)
        print("验证完成!")
        print("=" * 60)
        print(f"验证ID: {results['validation_id']}")
        print(f"验证分数: {results['validation_score']}%")
        print(f"总体状态: {results['overall_status']}")
        print(f"检查项目: {len(results['checks_performed'])}")
        print(f"真实性指标: {len(results['authenticity_indicators'])}")
        
        if results['anomalies_detected']:
            print(f"检测到异常: {len(results['anomalies_detected'])}")
        
        print(f"\n验证报告:")
        print(f"- JSON报告: {reports['json_report']}")
        print(f"- Markdown报告: {reports['markdown_report']}")
        
        return True
        
    except Exception as e:
        print(f"\n[!] 验证过程中发生错误: {str(e)}")
        return False

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)